'''
Author: tomwoo tom.woo@outlook.com
Date: 2025-07-15 00:01:02
LastEditors: tomwoo tom.woo@outlook.com
LastEditTime: 2025-07-22 17:57:22
FilePath: /multi-modal_agents/OCR_Pipelines.py
Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
'''

from unstructured.partition.pdf import partition_pdf
from unstructured.documents.elements import Text, Image, Table, CompositeElement


outputs_dir = "./outputs"


class PdfExtractionPipelines:
    def __init__(self, pdf_filename):
        self.pdf_filename = pdf_filename

    def load_and_split_document_for_text(self):
        print("processing document: ", self.pdf_filename)
        granular_elements = partition_pdf(
            self.pdf_filename,
            chunking_strategy="by_title",
            max_characters=2000,
            new_after_n_chars=1800,
            combine_text_under_n_chars=1000,
        )

        text_elements = [e.text.strip().replace("\n\n", "\n") for e in granular_elements if type(e) == Text or CompositeElement]
        print("Number of Detected Text Elements:", len(text_elements))

        return text_elements

    def load_and_split_document_for_tables(self):
        print("processing document: ", self.pdf_filename)
        raw_pdf_elements = partition_pdf(
            self.pdf_filename, 
            infer_table_structure=True, 
            strategy='hi_res'
        )

        table_elements = [e.metadata.text_as_html for e in raw_pdf_elements if type(e) == Table]
        print("Number of Detected HTML Table Elements:", len(table_elements))

        return table_elements

    def load_and_split_document_for_images(self):
        print("processing document: ", self.pdf_filename)
        image_text_elements = partition_pdf(
            self.pdf_filename,
            strategy="hi_res",
            hi_res_model_name="yolox",
            extract_images_in_pdf=True,
            extract_image_block_types=["Image"],
            extract_image_block_to_payload=False,
            extract_image_block_output_dir=outputs_dir
        )

        image_elements = [e.metadata.image_path for e in image_text_elements if type(e) is Image]
        print("Number of Detected Image Elements:", len(image_elements))

        return image_elements


def load_and_split_pdf_document(filename, type):
    pipeline = PdfExtractionPipelines(filename)

    match type:
        case "Text":
            return pipeline.load_and_split_document_for_text()
        case "Tables":
            return pipeline.load_and_split_document_for_tables()
        case "Images":
            return pipeline.load_and_split_document_for_images()
        case _:
            return None

# end of file
